import json

# protocol
# https://www.cvmart.net/topList/10098?tab=RankDescription
import os
import sys
import time

import cv2
import numpy as np

sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../../")

from recog.inference_crnn import get_crnn_model, get_lp_str

from PIL import Image, ImageDraw, ImageFont


def putTextPIL(im, str, loc, font, size, color, *args):
    im = Image.fromarray(im)
    font_path = os.path.dirname(os.path.abspath(__file__))
    ImageDraw.Draw(im).text(loc, str, color, font=ImageFont.truetype(font_path + "/simhei.ttf", int(15 * size)))
    im = np.array(im)
    return im


def init():
    """Initialize model
    Returns: model
    """
    cat_det_model = create_cat_det_model()
    lp_det_model = load_lp_det_model()
    crnn_model = get_crnn_model()
    return {
        "cat_det_model": cat_det_model,
        "lp_det_model": lp_det_model,
        "crnn_model": crnn_model,
    }


def crop_img(input_image, x1, y1, x2, y2):
    xc = (x1 + x2) / 2
    yc = (y1 + y2) / 2
    w = x2 - x1
    h = y2 - y1
    x_s = max(int(xc - w / 2), 0)
    x_e = min(input_image.shape[1], int(xc + w / 2))
    y_s = max(int(yc - h / 2), 0)
    y_e = min(input_image.shape[0], int(yc + h / 2))
    img_croped = input_image[y_s:y_e, x_s:x_e, :]
    return img_croped


def create_out_json(all_out_cars):
    root = {}

    is_alert = False
    target_count = 0
    root["algorithm_data"] = {}

    result = root["algorithm_data"]
    result['is_alert'] = is_alert
    result['target_count'] = 0
    if result.get('target_info') is None:
        result['target_info'] = []

    for a_car in all_out_cars:
        x1, y1, x2, y2 = a_car.car_coor
        if a_car.ocr is not None:
            result = root["algorithm_data"]
            is_alert = True
            result['is_alert'] = is_alert
            target_count += 1
            result['target_count'] = target_count

            result['target_info'].append({
                "points": [int(a) for a in list(a_car.lp_pts.reshape(-1))],
                'confidence': a_car.lp_pts_prob,
                "name": "plate",
                "color": a_car.lp_color,
                "ocr": a_car.ocr,
            })

    root['model_data'] = {}
    result = root['model_data']
    if result.get('objects') is None:
        result['objects'] = []

    for a_car in all_out_cars:
        if a_car.ocr is None:
            result['objects'].append({
                'x': a_car.car_coor[0],
                'y': a_car.car_coor[1],
                'height': a_car.car_coor[3] - a_car.car_coor[1],
                'width': a_car.car_coor[2] - a_car.car_coor[0],
                'confidence': a_car.car_prob,
                'name': a_car.cls_name,

            })
        else:
            result['objects'].append({
                'x': a_car.car_coor[0],
                'y': a_car.car_coor[1],
                'height': a_car.car_coor[3] - a_car.car_coor[1],
                'width': a_car.car_coor[2] - a_car.car_coor[0],
                'confidence': a_car.car_prob,
                'name': a_car.cls_name,

            })
            result['objects'].append({
                "points": [int(a) for a in list(a_car.lp_pts.reshape(-1))],
                'confidence': a_car.lp_pts_prob,
                "name": "plate",
                "color": a_car.lp_color,
                "ocr": a_car.ocr,
            })
    return root


def process_image(handle=None, input_image=None, args=None, **kwargs):
    """Do inference to analysis input_image and get output
    Attributes:
    handle: algorithm handle returned by init()
    input_image (numpy.ndarray): image to be process, format: (h, w, c), BGR
    Returns: process result
    """
    visualise = False
    all_out_cars = []
    t = time.time()
    boxes_prob_id_name = get_cat_boxes(input_image, handle['cat_det_model'])

    for x1, y1, x2, y2, prob, class_id, class_name in boxes_prob_id_name:
        # if class_name in ['car', 'bus', 'truck']:
        if True:
            croppd_im = crop_img(input_image, x1, y1, x2, y2)

            a_car = ScopeCarStru("None")
            a_car.car_coor = x1, y1, x2, y2
            a_car.car_prob = float(prob)
            # a_car.cls_name = class_name
            a_car.cls_name = "car"
            a_car.cls_ID = class_id

            # crop lp and recog
            if min(croppd_im.shape[0], croppd_im.shape[1]) > 200:
                coord, lp_prob, Ilp = get_lp_boxes(croppd_im, handle['lp_det_model'])
                if coord is not None:
                    coord = coord + np.array([x1, y1])
                    lp_str, color = get_lp_str(Ilp, handle['crnn_model'])
                    a_car.ocr = lp_str
                    a_car.lp_pts = coord
                    a_car.lp_color = color
                    a_car.lp_pts_prob = float(lp_prob)
                    a_car.warpped_im = Ilp

            all_out_cars.append(a_car)

    if visualise:
        for i, a_car in enumerate(all_out_cars):
            x1, y1, x2, y2 = a_car.car_coor
            cv2.rectangle(input_image, (x1, y1), (x2, y2), (0, 0, 255), 2)
            class_name = a_car.cls_name
            prob = a_car.car_prob
            cv2.putText(input_image, "{}:{:.2f}".format(class_name, float(prob)), (x1, y1),
                        cv2.FONT_HERSHEY_PLAIN, 2, (255, 255, 0), 2)

            if a_car.lp_pts is not None:
                cv2.polylines(input_image, [a_car.lp_pts.reshape(-1, 1, 2)], True, (255, 255, 0), 2)
                # cv2.putText(input_image, a_car.ocr, tuple(a_car.lp_pts[0]), cv2.FONT_HERSHEY_PLAIN, 2, (255, 0, 0), 2)
                input_image = putTextPIL(input_image, a_car.ocr,
                                         tuple(a_car.lp_pts[2]), cv2.FONT_HERSHEY_PLAIN, 2, (255, 255, 0), 2)
                Ilp = a_car.warpped_im
                cv2.imshow("Ilp{}".format(i), Ilp)

        cv2.namedWindow("x", cv2.WINDOW_NORMAL)
        cv2.imshow("x", input_image)
        cv2.resizeWindow("x", 1440, 900)

        cv2.waitKey(0)
        for i, a_car in enumerate(all_out_cars):
            cv2.destroyWindow("Ilp{}".format(i))

    # return boxes_prob_id_name, all_lp_strs
    print("processing img cost:{}".format(time.time() - t))
    json_root = create_out_json(all_out_cars)
    return json.dumps(json_root, indent=4, ensure_ascii=False)


# https://www.cvmart.net/topList/10098?dbType=1-3-2&tab=RankDescription
# # Process image here
# fake_result = {
#     'objects': []
# }
# fake_result['objects'].append({
#     'points': [x1, y1, x2, y2, ......, xn, yn],  # 点之间没有顺序关系
#     'name': 'plate',
#     'confidence': 0.5,
#     'ocr': '鲁B60176,
#            'color': 'blue'
# })
# return json.dumps(fake_result, indent=4)
#
def pad_bgr(bgr_im):
    offset = 40
    full_im = np.zeros([1920, 1920, 3], dtype=np.uint8)
    full_im[offset:bgr_im.shape[0] + offset, offset:bgr_im.shape[1] + offset, :] = bgr_im
    return full_im


if __name__ == '__main__':
    handle = init()

    root = "/home/leo/Downloads/data_set/license_plate_detect_set/val/"

    whole_num = 0
    correct_num = 0
    for file in os.listdir(root):
        if ".jpg" in file:
            bgr_im = cv2.imread(os.path.join(root, file))
            lp_str = file.split("_")[-2]
            bgr_im = pad_bgr(bgr_im)
            # cv2.imshow("ful", bgr_im)
            # cv2.waitKey(0)

            out_json = process_image(handle, bgr_im)
            # print(out_json)
            if json.loads(out_json)['algorithm_data']['is_alert'] is True:
                correct_num += (lp_str == json.loads(out_json)['algorithm_data']['target_info'][0]['ocr'])
            whole_num += 1
            print("acc:{}".format(correct_num / whole_num))
